import torch
import torch.nn as nn
import torchvision
import torchvision.transforms as transforms
from VGGConvNet import VGGConvNet
import matplotlib.pyplot as plt
import os

# Device configuration
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

classes = ('plane', 'car', 'bird', 'cat',
           'deer', 'dog', 'frog', 'horse', 'ship', 'truck')

transforms_train = transforms.Compose([
    # 以概率P对图片进行翻转，默认值为0.5
    transforms.RandomHorizontalFlip(),
    # transforms.ToTensor()将numpy的ndarray或PIL.Image读的图片转换成形状为(C,H, W)的Tensor格式，
    # 且/255归一化到[0,1.0]之间
    transforms.ToTensor(),
    transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))
])
transforms_test = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))
])

# Hyper parameters
num_epochs = 50
num_classes = 10
batch_size = 128
learning_rate = 0.0001
step_size = 10  # 每n个epoch更新一次学习率
log_dir = "./resources/ckpt/cifar_model.pth"

# Cifar10 dataset

'''
transforms.ToTensor()将numpy的ndarray或PIL.Image读的图片转换成形状为(C,H, W)的Tensor格式，
且/255归一化到[0,1.0]之间
'''
train_dataset = torchvision.datasets.CIFAR10(root='./resources/',
                                             train=True,
                                             transform=transforms_train,
                                             download=True)

test_dataset = torchvision.datasets.CIFAR10(root='./resources/',
                                            train=False,
                                            transform=transforms_test)

# Data loader
# 从已有的数据读取接口的输入按照batch size封装成Tensor,shuffle=True为每迭代一次洗牌重新按照batch_size分
train_loader = torch.utils.data.DataLoader(dataset=train_dataset,
                                           batch_size=batch_size,
                                           shuffle=True)

test_loader = torch.utils.data.DataLoader(dataset=test_dataset,
                                          batch_size=batch_size,
                                          shuffle=False)

model = VGGConvNet(num_classes).to(device)

# Loss and optimizer，交叉熵损失函数，Adam是比较好的优化器
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
# 更新学习率，按自己预设的值step_size，每step_size个epoch后学习率进行伽马衰减。注意这个要在epoch内循环的后面，否则训练时学习率会提前一个epoch更新
# scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=step_size, gamma=0.5, last_epoch=-1)

# Train the model,总的批数
total_step = len(train_loader)


if os.path.exists(log_dir):
    checkpoint = torch.load(log_dir)
    model.load_state_dict(checkpoint['model'])
    start_epoch = checkpoint['epoch']
    print("加载epoch成功".format(start_epoch))
else:
    start_epoch = 0
    print("无保存模型，从头开始训练")


train_total = 0
train_correct = 0
trainLoss_list = []
for epoch in range(start_epoch, num_epochs):
    for i, (images, labels) in enumerate(train_loader):
        images = images.to(device)
        labels = labels.to(device)

        # Forward pass
        outputs = model(images)
        loss = criterion(outputs, labels)

        trainLoss_list.append(loss.item())

        # 第一个值是具体的value（我们用下划线_表示），第二个值是value所在的index（也就是predicted）,1为所在行的最大值
        _, predicted = torch.max(outputs.data, 1)
        train_total += labels.size(0)
        train_correct += (predicted == labels).sum().item()

        # Backward and optimize，将梯度置零每一批次后，避免梯度叠加,因为使用的是当前参数空间对应的梯度
        optimizer.zero_grad()
        # 损失后向传播，计算梯度
        loss.backward()
        # 优化器更新参数
        optimizer.step()

        if (i + 1) % 100 == 0:
            print('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}'
                  .format(epoch + 1, num_epochs, i + 1, total_step, loss.item()))
            print('Epoch [{}/{}],Train Accuracy of the model on train images: {:.2f} %'
                  .format(epoch + 1, num_epochs, 100 * train_correct / train_total))
    # scheduler.step()
    state = {'model': model.state_dict(), 'epoch': epoch}
    torch.save(state, log_dir)


print("train accuracy:{:.2f} % ".format(100 * train_correct / train_total))

# Test the model
'''
model.eval()是保证BN层能够用全部训练数据的均值和方差，即测试过程中要保证BN层的均值和方差不变。
对于Dropout，model.eval()是利用到了所有网络连接，即不进行随机舍弃神经元。
'''

class_correct = list(0. for i in range(num_classes))
class_total = list(0. for i in range(num_classes))


model.eval()  # eval mode (batchnorm uses moving mean/variance instead of mini-batch mean/variance)
# 强制之后的内容不进行计算图构建
with torch.no_grad():
    correct = 0
    total = 0
    for images, labels in test_loader:
        images = images.to(device)
        labels = labels.to(device)
        outputs = model(images)
        # 第一个值是具体的value（我们用下划线_表示），第二个值是value所在的index（也就是predicted）,1为所在行的最大值
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

        c = (predicted == labels).squeeze()

        for i in range(len(labels)):
            _label = labels[i]
            class_correct[_label] += c[i].item()
            class_total[_label] += 1

    print('Test Accuracy of the model on the 10000 test images: {} %'.format(100 * correct / total))

    for i in range(num_classes):
        print('Accuracy of %5s : %2d %%' % (
            classes[i], 100 * class_correct[i] / class_total[i]))

plt.plot(trainLoss_list)
plt.show()

